package com.itcam.ds_.balancetree;

/**
 * Author:Cammy.Wu
 * Date:2024-06-06
 * Description:
 */

public class AVLTree {
    static class AVLNode {
        int key;
        Object value;
        AVLNode left;
        AVLNode right;
        int height = 1;

        public AVLNode(int key, Object value) {
            this.key = key;
            this.value = value;
        }

        public AVLNode(int key) {
            this.key = key;
        }

        public AVLNode(int key, Object value, AVLNode left, AVLNode right) {
            this.key = key;
            this.value = value;
            this.left = left;
            this.right = right;
        }

        // 求节点高度
        private int height(AVLNode node) {
            return node == null ? 0 : node.height;
        }

        // 更新节点高度
        private void updateNodeHeight(AVLNode node) {
            node.height = Integer.max(height(node.left), height(node.right)) + 1;
        }

        // 平衡因子（balance factor） = 左子树高度-右子树高度
        private int balanceFactor(AVLNode node) {
            /*
             * 0,1,-1 平衡
             *  balanceFactor > 1 左高
             *  balanceFactor < -1 右高
             */
            return height(node.left) - height(node.right);
        }

        /*
            LL
                -失衡节点的 balanceFactor > 1
                -失衡节点的左孩子的 balanceFactor >= 0
                -★ 左边高，且左孩子这边也是左边高
            LR
                -失衡节点的 balanceFactor > 1
                -失衡节点的左孩子的 balanceFactor < 0
                -★ 左边高，且左孩子这边是右边高
            RL
                -失衡节点的 balanceFactor < -1
                -失衡节点的右孩子的 balanceFactor >= 0
                -★ 右边高，且右孩子这边是左边高
            RR
                -失衡节点的 balanceFactor < -1
                -失衡节点的右孩子的 balanceFactor < 0
                -★ 右边高，且右孩子这边也是右边高
         */

        // LL
        // 右旋
        private AVLNode rightRotate(AVLNode parent) { // 参数：要旋转的节点 返回值：新的根节点
            AVLNode left = parent.left;
            AVLNode leftRight = left.right;
            left.right = parent;
            parent.left = leftRight;
            updateNodeHeight(parent);
            updateNodeHeight(left);
            return left;
        }

        // RR
        // 左旋
        private AVLNode leftRotate(AVLNode parent) {
            AVLNode right = parent.right;
            AVLNode rightLeft = right.left;
            right.left = parent;
            parent.right = rightLeft;
            updateNodeHeight(parent);
            updateNodeHeight(right);
            return right;
        }

        // LR
        // 先左旋左子树，再右旋根节点
        private AVLNode leftRightRotate(AVLNode parent) {
            parent.left = leftRotate(parent.left);
            return rightRotate(parent);
        }

        // RL
        // 先右旋右子树，再左旋根节点
        private AVLNode rightLeftRotate(AVLNode parent) {
            parent.right = rightRotate(parent.right);
            return leftRotate(parent);
        }

        // 检查节点是否失衡，重新平衡代码
        private AVLNode balance(AVLNode node) {
            if (node == null) {
                return null;
            }
            int bf = balanceFactor(node);
            if (bf > 1 && balanceFactor(node.left) >= 0) {
                // LL
                return rightRotate(node);
            } else if (bf > 1 && balanceFactor(node.left) < 0) {
                // LR
                return leftRightRotate(node);
            } else if (bf < -1 && balanceFactor(node.right) > 0) {
                // RL
                return rightLeftRotate(node);
            } else if (bf < -1 && balanceFactor(node.right) <= 0) {
                // RR
                return leftRotate(node);
            }
            return node;
        }

        AVLNode root;

        // 添加新节点
        public void put(int key, Object value) {
            root = doPut(root, key, value);
        }

        // 创建新节点
        private AVLNode doPut(AVLNode node, int key, Object value) {
            // 1.找到空位，创建新节点
            if (node == null) {
                return new AVLNode(key, value);
            }
            // 2.key存在，更新
            if (node.key == key) {
                node.value = value;
                return node;
            }
            // 3.继续查找
            if (key < node.key) {
                node.left = doPut(node.left, key, value);
            } else {
                node.right = doPut(node.right, key, value);
            }

            updateNodeHeight(node);
            return balance(node);
        }

        // 删除节点
        public void remove(int key) {
            root = doRemove(root, key);
        }

        // 找打待删除节点
        private AVLNode doRemove(AVLNode node, int key) {
            // 1.node == null
            if (node == null) {
                return null;
            }
            // 2.没找到key
            if (key < node.key) {
                node.left = doRemove(node.left, key);
            } else if (key > node.key) {
                node.right = doRemove(node.right, key);
            } else {
                // 3.找到key 1)没有孩子 2)有一个孩子 3)有两个孩子
                if (node.left == null && node.right == null) {
                    // 1)没有孩子
                    return null;
                } else if (node.left == null) {
                    // 2)有一个孩子
                    node = node.right;
                } else if (node.right == null) {
                    // 2)有一个孩子
                    node = node.left;
                } else {
                    // todo:3)有两个孩子
                    AVLNode nodeRight = node.right;
                    while (nodeRight.left != null) {
                        nodeRight = nodeRight.left;
                    }
                    nodeRight.right = doRemove(node.right, nodeRight.key);
                    nodeRight.left = node.left;
                    node = nodeRight;
                }
            }
            // 4.更新高度
            updateNodeHeight(node);
            // 5.balance
            return balance(node);
        }
    }
}
